package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

var (
	dbPath        = "stock_data.db"
	sleepDuration = 200 * time.Millisecond
)

var DebugFlag = false

func main() {

	flag.StringVar(&dbPath, "d", "stock_data.db", "sqlite3 save data file,default stock_data.db")

	db, err := initDB()
	if err != nil {
		fmt.Printf("初始化数据库失败: %v\n", err)
		return
	}
	defer db.Close()

	startTime := time.Now().UnixMilli()

	fmt.Println("开始同步股票数据...")
	if err := syncStocks(db); err != nil {
		fmt.Printf("同步股票数据失败: %v\n", err)
	}
	endTime := time.Now().UnixMilli()
	fmt.Println("use time:ms", (endTime - startTime))
	startTime = endTime

	fmt.Println("开始同步概念板块数据...")
	if err := syncConcepts(db); err != nil {
		fmt.Printf("同步概念板块数据失败: %v\n", err)
	}
	endTime = time.Now().UnixMilli()
	fmt.Println("use time:ms", (endTime - startTime))
	startTime = endTime

	fmt.Println("开始同步股票概念关联数据...")
	if err := syncStockConcepts(db); err != nil {
		fmt.Printf("同步股票概念关联数据失败: %v\n", err)
	}
	endTime = time.Now().UnixMilli()
	fmt.Println("use time:ms", (endTime - startTime))
	startTime = endTime

	fmt.Println("开始同步概念龙头股数据...")
	if err := syncConceptStocks(db); err != nil {
		fmt.Printf("同步概念龙头股数据失败: %v\n", err)
	}
	endTime = time.Now().UnixMilli()
	fmt.Println("use time:ms", (endTime - startTime))
	startTime = endTime
	fmt.Println("数据同步完成")
}

func initDB() (*sql.DB, error) {
	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, err
	}

	// 创建股票表
	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS stocks (
		code TEXT PRIMARY KEY,
		name TEXT,
		create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
		delete_time DATETIME,
		is_deleted INTEGER DEFAULT 0
	)`)
	if err != nil {
		return nil, err
	}

	// 创建概念板块表
	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS concepts (
		code TEXT PRIMARY KEY,
		name TEXT,
		create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
		delete_time DATETIME,
		is_deleted INTEGER DEFAULT 0
	)`)
	if err != nil {
		return nil, err
	}

	// 创建股票概念关联表
	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS stock_concepts (
		stock_code TEXT,
		concept_code TEXT,
		reason TEXT,
		create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
		delete_time DATETIME,
		is_deleted INTEGER DEFAULT 0,
		PRIMARY KEY (stock_code, concept_code)
	)`)
	if err != nil {
		return nil, err
	}

	// 创建概念龙头股表
	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS concept_stocks (
		concept_code TEXT,
		concept_name TEXT,
		stock_code TEXT,
		stock_name TEXT,
		create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
		delete_time DATETIME,
		is_deleted INTEGER DEFAULT 0,
		PRIMARY KEY (concept_code, stock_code)
	)`)
	if err != nil {
		return nil, err
	}

	return db, nil
}

type StockResponse struct {
	Data struct {
		Diff []struct {
			Code string `json:"f12"`
			Name string `json:"f14"`
		} `json:"diff"`
	} `json:"data"`
}

func syncStocks(db *sql.DB) error {
	time.Sleep(sleepDuration)

	// 分页获取所有股票数据
	page := 1
	pageSize := 500
	allStocks := make([]struct {
		Code string `json:"f12"`
		Name string `json:"f14"`
	}, 0)

	for {
		time.Sleep(sleepDuration)

		params := url.Values{}
		params.Set("pn", fmt.Sprintf("%d", page))
		params.Set("pz", fmt.Sprintf("%d", pageSize))
		params.Set("po", "1")
		params.Set("np", "1")
		params.Set("ut", "bd1d9ddb04089700cf9c27f6f7426281")
		params.Set("fltt", "2")
		params.Set("invt", "2")
		params.Set("fid", "f3")
		params.Set("fs", "m:0+t:6,m:0+t:13,m:0+t:80,m:1+t:2,m:1+t:23")
		params.Set("fields", "f12,f14")

		var resp StockResponse
		if err := httpGet("https://20.push2.eastmoney.com/api/qt/clist/get", params, nil, &resp); err != nil {
			return fmt.Errorf("获取股票数据失败(第%d页): %v", page, err)
		}

		if len(resp.Data.Diff) == 0 {
			break
		}

		allStocks = append(allStocks, resp.Data.Diff...)
		page++
	}

	// 获取数据库中现有股票
	existingStocks := make(map[string]bool)
	rows, err := db.Query("SELECT code FROM stocks WHERE is_deleted = 0")
	if err != nil {
		return fmt.Errorf("查询现有股票失败: %v", err)
	}
	defer rows.Close()

	for rows.Next() {
		var code string
		if err := rows.Scan(&code); err != nil {
			return fmt.Errorf("扫描股票代码失败: %v", err)
		}
		existingStocks[code] = true
	}

	fmt.Println("req size:", len(allStocks), "db size:", len(existingStocks))

	// 处理新增和更新的股票
	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("开启事务失败: %v", err)
	}

	for _, stock := range allStocks {
		if _, exists := existingStocks[stock.Code]; exists {
			// 更新股票名称
			_, err := tx.Exec("UPDATE stocks SET name = ? WHERE code = ?", stock.Name, stock.Code)
			if err != nil {
				tx.Rollback()
				return fmt.Errorf("更新股票失败: %v", err)
			}
			delete(existingStocks, stock.Code)
		} else {
			// 新增股票
			_, err := tx.Exec("INSERT INTO stocks (code, name) VALUES (?, ?)", stock.Code, stock.Name)
			if err != nil {
				tx.Rollback()
				return fmt.Errorf("插入股票失败: %v", err)
			}
		}
	}

	// 标记数据库中不再存在的股票为已删除
	for code := range existingStocks {
		_, err := tx.Exec("UPDATE stocks SET is_deleted = 1, delete_time = CURRENT_TIMESTAMP WHERE code = ?", code)
		if err != nil {
			tx.Rollback()
			return fmt.Errorf("标记删除股票失败: %v", err)
		}
	}

	return tx.Commit()
}

type ConceptResponse struct {
	Data struct {
		Diff []struct {
			Code string `json:"f12"`
			Name string `json:"f14"`
		} `json:"diff"`
	} `json:"data"`
}

func syncConcepts(db *sql.DB) error {
	time.Sleep(sleepDuration)

	params := url.Values{}
	params.Set("key", "f62")
	params.Set("code", "m:90+t:3")

	var resp ConceptResponse
	if err := httpGet("https://data.eastmoney.com/dataapi/bkzj/getbkzj", params, nil, &resp); err != nil {
		return fmt.Errorf("获取概念板块数据失败: %v", err)
	}

	// 获取数据库中现有概念板块
	existingConcepts := make(map[string]bool)
	rows, err := db.Query("SELECT code FROM concepts WHERE is_deleted = 0")
	if err != nil {
		return fmt.Errorf("查询现有概念板块失败: %v", err)
	}
	defer rows.Close()

	for rows.Next() {
		var code string
		if err := rows.Scan(&code); err != nil {
			return fmt.Errorf("扫描概念板块代码失败: %v", err)
		}
		existingConcepts[code] = true
	}

	// 处理新增和更新的概念板块
	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("开启事务失败: %v", err)
	}

	for _, concept := range resp.Data.Diff {
		if _, exists := existingConcepts[concept.Code]; exists {
			// 更新概念板块名称
			_, err := tx.Exec("UPDATE concepts SET name = ? WHERE code = ?", concept.Name, concept.Code)
			if err != nil {
				tx.Rollback()
				return fmt.Errorf("更新概念板块失败: %v", err)
			}
			delete(existingConcepts, concept.Code)
		} else {
			// 新增概念板块
			_, err := tx.Exec("INSERT INTO concepts (code, name) VALUES (?, ?)", concept.Code, concept.Name)
			if err != nil {
				tx.Rollback()
				return fmt.Errorf("插入概念板块失败: %v", err)
			}
		}
	}

	// 标记数据库中不再存在的概念板块为已删除
	for code := range existingConcepts {
		_, err := tx.Exec("UPDATE concepts SET is_deleted = 1, delete_time = CURRENT_TIMESTAMP WHERE code = ?", code)
		if err != nil {
			tx.Rollback()
			return fmt.Errorf("标记删除概念板块失败: %v", err)
		}
	}

	return tx.Commit()
}

type StockConceptResponse struct {
	Result struct {
		Data []struct {
			StockCode   string `json:"SECURITY_CODE"`
			ConceptCode string `json:"NEW_BOARD_CODE"`
			ConceptName string `json:"BOARD_NAME"`
			Reason      string `json:"SELECTED_BOARD_REASON"`
		} `json:"data"`
	} `json:"result"`
}

func syncStockConcepts(db *sql.DB) error {
	// 获取所有股票代码
	rows, err := db.Query("SELECT code FROM stocks WHERE is_deleted = 0")
	if err != nil {
		return fmt.Errorf("查询股票代码失败: %v", err)
	}
	defer rows.Close()

	var stockCodes []string
	for rows.Next() {
		var code string
		if err := rows.Scan(&code); err != nil {
			return fmt.Errorf("扫描股票代码失败: %v", err)
		}
		stockCodes = append(stockCodes, code)
	}

	// 遍历所有股票查询概念板块
	for i, code := range stockCodes {
		if DebugFlag && i > 2 {
			break
		}
		time.Sleep(sleepDuration)

		// 根据股票代码前缀判断市场
		market := ".SH"
		if code[0] == '3' || code[0] == '0' {
			market = ".SZ"
		}
		params := url.Values{}
		params.Set("reportName", "RPT_F10_CORETHEME_BOARDTYPE")
		params.Set("columns", "SECUCODE,SECURITY_CODE,SECURITY_NAME_ABBR,NEW_BOARD_CODE,BOARD_NAME,SELECTED_BOARD_REASON,IS_PRECISE,BOARD_RANK,BOARD_YIELD,DERIVE_BOARD_CODE")
		params.Set("filter", fmt.Sprintf("(SECUCODE=\"%s\")(IS_PRECISE=\"1\")", code+market))
		params.Set("pageNumber", "1")
		params.Set("pageSize", "100")

		var resp StockConceptResponse
		//https://datacenter.eastmoney.com/securities/api/data/get
		//https://datacenter-web.eastmoney.com/api/data/v1/get
		if err := httpGet("https://datacenter-web.eastmoney.com/api/data/v1/get", params, nil, &resp); err != nil {
			fmt.Printf("获取股票%s概念板块失败: %v\n", code, err)
			continue
		}

		if len(resp.Result.Data) == 0 {
			fmt.Printf("股票%s无概念板块数据\n%v\n", code, resp)
			continue
		}

		// 获取该股票现有的概念关联
		existingConcepts := make(map[string]bool)
		rows, err := db.Query("SELECT concept_code FROM stock_concepts WHERE stock_code = ? AND is_deleted = 0", code)
		if err != nil {
			fmt.Printf("查询股票%s现有概念关联失败: %v\n", code, err)
			continue
		}
		defer rows.Close()

		for rows.Next() {
			var conceptCode string
			if err := rows.Scan(&conceptCode); err != nil {
				fmt.Printf("扫描股票%s概念代码失败: %v\n", code, err)
				continue
			}
			existingConcepts[conceptCode] = true
		}

		// 处理新增和更新的概念关联
		tx, err := db.Begin()
		if err != nil {
			fmt.Printf("开启事务失败: %v\n", err)
			continue
		}

		for _, item := range resp.Result.Data {
			if _, exists := existingConcepts[item.ConceptCode]; exists {
				// 更新概念关联信息
				_, err := tx.Exec("UPDATE stock_concepts SET reason = ? WHERE stock_code = ? AND concept_code = ?",
					item.Reason, item.StockCode, item.ConceptCode)
				if err != nil {
					tx.Rollback()
					fmt.Printf("更新股票%s概念关联失败: %v\n", code, err)
					continue
				}
				delete(existingConcepts, item.ConceptCode)
			} else {
				// 新增概念关联
				_, err := tx.Exec("INSERT INTO stock_concepts (stock_code, concept_code, reason) VALUES (?, ?, ?)",
					item.StockCode, item.ConceptCode, item.Reason)
				if err != nil {
					tx.Rollback()
					fmt.Printf("插入股票%s概念关联失败: %v\n", code, err)
					continue
				}
			}
		}

		// 标记数据库中不再存在的概念关联为已删除
		for conceptCode := range existingConcepts {
			_, err := tx.Exec("UPDATE stock_concepts SET is_deleted = 1, delete_time = CURRENT_TIMESTAMP WHERE stock_code = ? AND concept_code = ?",
				code, conceptCode)
			if err != nil {
				tx.Rollback()
				fmt.Printf("标记删除股票%s概念关联失败: %v\n", code, err)
				continue
			}
		}

		if err := tx.Commit(); err != nil {
			fmt.Printf("提交事务失败: %v\n", err)
		}
	}

	return nil
}

type ConceptStockResponse struct {
	Result []struct {
		StockCode string `json:"SECUCODE"`
		StockName string `json:"SECURITY_NAME_ABBR"`
	} `json:"result"`
}

func syncConceptStocks(db *sql.DB) error {
	// 获取所有概念板块代码
	rows, err := db.Query("SELECT code, name FROM concepts WHERE is_deleted = 0")
	if err != nil {
		return fmt.Errorf("查询概念板块代码失败: %v", err)
	}
	defer rows.Close()

	var concepts []struct {
		Code string
		Name string
	}
	for rows.Next() {
		var code, name string
		if err := rows.Scan(&code, &name); err != nil {
			return fmt.Errorf("扫描概念板块代码失败: %v", err)
		}
		concepts = append(concepts, struct {
			Code string
			Name string
		}{code, name})
	}

	// 遍历所有概念板块查询龙头股
	for i, concept := range concepts {
		if DebugFlag && i > 2 {
			break
		}

		time.Sleep(sleepDuration)

		params := url.Values{}
		params.Set("type", "RTP_F10_POPULAR_LEADING")
		params.Set("sty", "SECUCODE,SECURITY_NAME_ABBR") //SECUCODE,SECURITY_NAME_ABBR
		params.Set("extraCols", "f2~01~SECURITY_CODE~NEWEST_PRICE,f3~01~SECURITY_CODE~YIELD")
		params.Set("params", strings.ReplaceAll(concept.Code, "BK", "BI"))
		params.Set("sr", "")
		params.Set("st", "")
		params.Set("source", "HSF10")
		params.Set("client", "PC")
		params.Set("v", "06791922568983983")

		var resp ConceptStockResponse
		if err := httpGet("https://datacenter.eastmoney.com/securities/api/data/get", params, nil, &resp); err != nil {
			fmt.Printf("获取概念板块%s龙头股失败: %v\n", concept.Code, err)
			continue
		}

		// 获取该概念现有的龙头股
		existingStocks := make(map[string]bool)
		rows, err := db.Query("SELECT stock_code FROM concept_stocks WHERE concept_code = ? AND is_deleted = 0", concept.Code)
		if err != nil {
			fmt.Printf("查询概念板块%s现有龙头股失败: %v\n", concept.Code, err)
			continue
		}
		defer rows.Close()

		for rows.Next() {
			var stockCode string
			if err := rows.Scan(&stockCode); err != nil {
				fmt.Printf("扫描概念板块%s龙头股代码失败: %v\n", concept.Code, err)
				continue
			}
			existingStocks[stockCode] = true
		}

		// 处理新增和更新的龙头股
		tx, err := db.Begin()
		if err != nil {
			fmt.Printf("开启事务失败: %v\n", err)
			continue
		}

		for _, item := range resp.Result {
			if _, exists := existingStocks[item.StockCode]; exists {
				// 更新龙头股信息
				_, err := tx.Exec("UPDATE concept_stocks SET stock_name = ? WHERE concept_code = ? AND stock_code = ?",
					item.StockName, concept.Code, item.StockCode)
				if err != nil {
					tx.Rollback()
					fmt.Printf("更新概念板块%s龙头股失败: %v\n", concept.Code, err)
					continue
				}
				delete(existingStocks, item.StockCode)
			} else {
				// 新增龙头股
				_, err := tx.Exec("INSERT INTO concept_stocks (concept_code, concept_name, stock_code, stock_name) VALUES (?, ?, ?, ?)",
					concept.Code, concept.Name, item.StockCode, item.StockName)
				if err != nil {
					tx.Rollback()
					fmt.Printf("插入概念板块%s龙头股失败: %v\n", concept.Code, err)
					continue
				}
			}
		}

		// 标记数据库中不再存在的龙头股为已删除
		for stockCode := range existingStocks {
			_, err := tx.Exec("UPDATE concept_stocks SET is_deleted = 1, delete_time = CURRENT_TIMESTAMP WHERE concept_code = ? AND stock_code = ?",
				concept.Code, stockCode)
			if err != nil {
				tx.Rollback()
				fmt.Printf("标记删除概念板块%s龙头股失败: %v\n", concept.Code, err)
				continue
			}
		}

		if err := tx.Commit(); err != nil {
			fmt.Printf("提交事务失败: %v\n", err)
		}
	}

	return nil
}

func httpGet(baseURL string, params url.Values, headers http.Header, result interface{}) error {
	// 构建完整URL
	urlStr := baseURL
	if params != nil {
		urlStr += "?" + params.Encode()
	}

	client := &http.Client{}
	req, err := http.NewRequest("GET", urlStr, nil)
	if err != nil {
		return fmt.Errorf("创建请求失败: %v", err)
	}

	// 设置默认请求头
	defaultHeaders := http.Header{
		"User-Agent":       {"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"},
		"Referer":          {"https://data.eastmoney.com/"},
		"Accept":           {"application/json, text/javascript, */*; q=0.01"},
		"Accept-Language":  {"zh-CN,zh;q=0.9,en;q=0.8"},
		"X-Requested-With": {"XMLHttpRequest"},
	}

	// 合并默认请求头和自定义请求头
	if headers == nil {
		headers = make(http.Header)
	}
	for k, v := range defaultHeaders {
		if _, ok := headers[k]; !ok {
			headers[k] = v
		}
	}
	req.Header = headers

	resp, err := client.Do(req)

	if err != nil {
		fmt.Println(urlStr)
		return fmt.Errorf("请求失败: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		fmt.Println(urlStr)
		return fmt.Errorf("API返回非200状态码: %d", resp.StatusCode)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(urlStr)
		return fmt.Errorf("读取响应体失败: %v", err)
	}
	err = json.Unmarshal(body, result)
	if DebugFlag && err == nil {
		fmt.Println(urlStr, "\n", string(body))
	}

	return err
}
